{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dac8d8a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "import cv2\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c1b11b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def im_show(image):\n",
    "    plt.imshow(image)\n",
    "    plt.show()\n",
    "    \n",
    "def read_view_in(in_path):\n",
    "    frame_id_2_persons = {}\n",
    "    with open(in_path, 'r') as fin:\n",
    "        for line in fin.readlines():\n",
    "            items = line.split(',')\n",
    "            frame_id = int(items[0])\n",
    "            if frame_id in frame_id_2_persons:\n",
    "                frame_id_2_persons[frame_id].append(items[1:])\n",
    "            else:\n",
    "                frame_id_2_persons[frame_id] = [items[1:]]\n",
    "    return frame_id_2_persons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02dab5f0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#input_video = '/mnt/nas/04_自由视角处理后素材/球员识别测试164/view_004_test.ts'\n",
    "\n",
    "# input_video = 'view_out/029_to_004_4.mp4'\n",
    "# input_view_poin_path = 'view_out/034_to_004.txt'\n",
    "# output_video = 'view_out/034_to_004_5.mp4'\n",
    "\n",
    "input_video = 'view_out/034_to_004_5.mp4'\n",
    "input_view_poin_path = 'boxes/view_004_boxes.txt'\n",
    "output_video = 'view_out/0004_to_004_6_1.mp4'\n",
    "\n",
    "frame_id_2_persons = read_view_in(input_view_poin_path)\n",
    "frame_id = -1\n",
    "more = True\n",
    "\n",
    "reader = cv2.VideoCapture(input_video)\n",
    "width = int(reader.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "height = int(reader.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "#width, height = 1920, 1080\n",
    "print(width, height)\n",
    "writer = cv2.VideoWriter(output_video, cv2.VideoWriter_fourcc(*\"mp4v\"), 25, (width, height))\n",
    "while more:\n",
    "    more, frame = reader.read()\n",
    "    frame_id += 1\n",
    "    if(frame_id % 20 == 0):\n",
    "        print(\"== handle \", frame_id)\n",
    "    if(frame_id in frame_id_2_persons):\n",
    "        persons = frame_id_2_persons[frame_id]\n",
    "        for person in persons:\n",
    "            num_txt, x, y , width, height = person\n",
    "            x,y,height = int(x), int(y),int(height)\n",
    "            y = y+height\n",
    "            if(x>0 and y>0):\n",
    "                cv2.circle(frame, (x,y), 3, (255, 255, 255), 5)\n",
    "                #cv2.putText(frame, num_txt, (x,y), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 1)\n",
    "    writer.write(frame)\n",
    "    \n",
    "reader.release()\n",
    "writer.release()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dec3030",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "\n",
    "def read_valid_number(num_path):\n",
    "    valid_num = []\n",
    "    with open(num_path, 'r') as fin:\n",
    "        for line in fin.readlines():\n",
    "            valid_num.append(line.strip())\n",
    "    return valid_num\n",
    "\n",
    "#valid_nums = read_valid_number('./valid_num.txt')\n",
    "\n",
    "def point_in_rect(c_x, c_y, rect_x, rect_y, rect_height, rect_width):\n",
    "    rect_b = rect_y+rect_height\n",
    "    rect_r = rect_x+rect_width\n",
    "    \n",
    "    if (c_x >= rect_x) and (c_x <= rect_r) and (c_y >= rect_y) and (c_y <= rect_b):\n",
    "        return True\n",
    "    return False\n",
    "\n",
    "input_view_point_pathes = ['view_out/000_to_014_num_v2.txt', \n",
    "                           'view_out/025_to_014_num_v2.txt',\n",
    "                           'view_out/031_to_014_num_v2.txt']\n",
    "\n",
    "boxes_path = 'boxes/view_014_number_v2_all_boxes_2.txt'\n",
    "unknown_path = 'boxes/view_014_number_v2_unknown.txt'\n",
    "#input_video = 'view_out/0004_to_004_6.mp4'\n",
    "input_video = '/mnt/nas/03_自由视角原始素材/202303_ZCAM相机_CBA-东莞_23-24/广东vs浙江_fv2_test/view_014_test.mp4'\n",
    "output_video = \"view_out/view_4camera_demo_unknown.mp4\"\n",
    "\n",
    "frame_id = -1\n",
    "more = True\n",
    "frame_id_2_person_list = list(map(read_view_in, input_view_point_pathes))\n",
    "\n",
    "frame_id_2_unknown_box = read_view_in(unknown_path)\n",
    "\n",
    "frame_id_2_person_box = read_view_in(boxes_path)\n",
    "\n",
    "reader = cv2.VideoCapture(input_video)\n",
    "\n",
    "width = int(reader.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "height = int(reader.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "\n",
    "colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 0, 255)]#, (0, 255, 255), (255, 255, 255)]\n",
    "\n",
    "writer = cv2.VideoWriter(output_video, cv2.VideoWriter_fourcc(*\"mp4v\"), 25, (width, height))\n",
    "\n",
    "while more:\n",
    "    more, frame = reader.read()\n",
    "    frame_id += 1\n",
    "    #print(\"== read frame \", frame_id)\n",
    "    if(frame_id % 20 == 0):\n",
    "        print(\"== handle \", frame_id)\n",
    "        \n",
    "    person_2_boxes = {}\n",
    "    for idx, frame_id_2_persons in enumerate(frame_id_2_person_list):\n",
    "        if(frame_id in frame_id_2_persons):\n",
    "            #print(\"== frame id \", frame_id)\n",
    "            persons = frame_id_2_persons[frame_id]\n",
    "            for person in persons:\n",
    "                num_txt, x, y, p_width, p_height = person\n",
    "                #num_txt, x, y, b, r = person\n",
    "                x, y, p_width, p_height = int(x), int(y), int(p_width), int(p_height)\n",
    "                #x, y , b, r = int(x), int(y), int(b), int(r)\n",
    "                if(x>0 and x<width and y>0 and y<height):\n",
    "                    if(num_txt in person_2_boxes):\n",
    "                        person_2_boxes[num_txt].append((idx, x, y, p_width, p_height))\n",
    "                        #person_2_boxes[num_txt].append((idx, x, y, b, r))\n",
    "                    else:\n",
    "                        person_2_boxes[num_txt] = [(idx, x, y, p_width, p_height)]\n",
    "                        #person_2_boxes[num_txt] = [(idx, x, y, b, r)]\n",
    "        \n",
    "    if frame_id in frame_id_2_person_box:\n",
    "        view_0_persons = frame_id_2_person_box[frame_id]\n",
    "        for person in view_0_persons:\n",
    "            num_txt, p_x, p_y, p_width, p_height = person\n",
    "            p_x, p_y, p_width, p_height = int(p_x), int(p_y), int(p_width), int(p_height)\n",
    "            if num_txt in person_2_boxes:\n",
    "                person_2_boxes[num_txt].append((3, p_x, p_y, p_width, p_height))\n",
    "            else:\n",
    "                person_2_boxes[num_txt] = [(3, p_x, p_y, p_width, p_height)]\n",
    "            \n",
    "    draw_boxes = []\n",
    "    for person_num in person_2_boxes:\n",
    "        persons = person_2_boxes[person_num]\n",
    "        for person in persons:\n",
    "            idx, p_x, p_y, p_width, p_height = person\n",
    "            color = colors[idx]\n",
    "            #print(\"idx circle point \", (idx, p_x, p_y))\n",
    "            cv2.circle(frame, (p_x,p_y), 3, color, 5)\n",
    "            \n",
    "        person_last = persons[-1]\n",
    "        idx, p_x, p_y, p_width, p_height = person_last\n",
    "        if idx == 3:\n",
    "            p_r = p_x+p_width\n",
    "            p_b = p_y+p_height\n",
    "            cv2.putText(frame, person_num, (p_x, p_y), cv2.FONT_HERSHEY_COMPLEX, 1, (255, 255, 255), 1)\n",
    "            cv2.rectangle(frame, (p_x, p_y), (p_r, p_b), (255, 255, 255), 2)\n",
    "            draw_boxes.append((p_x, p_y))\n",
    "        else:\n",
    "            color = colors[idx]\n",
    "            #cv2.putText(frame, person_num, (p_x, p_y), cv2.FONT_HERSHEY_COMPLEX, 1, (255, 255, 255), 1)\n",
    "    \n",
    "    if frame_id in frame_id_2_unknown_box:\n",
    "        unknown_boxes = frame_id_2_unknown_box[frame_id]\n",
    "        for unknow_person in unknown_boxes:\n",
    "            p_num, p_x, p_y, p_width, p_height = unknow_person\n",
    "            p_x, p_y, p_width, p_height = int(p_x), int(p_y), int(p_width), int(p_height)\n",
    "            p_r = p_x + p_width\n",
    "            p_b = p_y + p_height\n",
    "            if((p_x, p_y) in draw_boxes):\n",
    "                continue\n",
    "            \n",
    "            cv2.rectangle(frame, (p_x, p_y), (p_r, p_b), (255, 0, 0), 2)\n",
    "            \n",
    "            hit_person_circles = []\n",
    "            for person_num in person_2_boxes:\n",
    "                persons = person_2_boxes[person_num]\n",
    "                for person in persons:\n",
    "                    idx, c_x, c_y, c_width, c_height = person\n",
    "                    if (idx == 3):# ignore current view\n",
    "                        continue \n",
    "                    if point_in_rect(c_x, c_y, p_x, p_y, p_height, p_width):\n",
    "                        hit_person_circles.append(person_num)\n",
    "                        \n",
    "            if len(hit_person_circles)>0:\n",
    "                person_num = hit_person_circles[0]\n",
    "                cv2.putText(frame, person_num, (p_x, p_y), cv2.FONT_HERSHEY_COMPLEX, 1, (255, 0, 0), 1)\n",
    " \n",
    "    \n",
    "    writer.write(frame)\n",
    "    \n",
    "reader.release()\n",
    "writer.release()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c78566d4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
